#!/usr/bin/env python3
import rospy
import heapq
import numpy as np
from nav_msgs.msg import OccupancyGrid, Path
from geometry_msgs.msg import PoseStamped

class Node:
    """A* 节点"""
    def __init__(self, position, parent=None):
        self.position = position
        self.parent = parent
        self.g = 0
        self.h = 0
        self.f = 0

    def __lt__(self, other):
        return self.f < other.f

def heuristic(a, b):
    """曼哈顿距离启发函数"""
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def astar(grid_map, start, end):
    """A*路径规划"""
    open_list = []
    closed_set = set()
    start_node = Node(start)
    end_node = Node(end)
    heapq.heappush(open_list, start_node)

    while open_list:
        current_node = heapq.heappop(open_list)

        if current_node.position == end_node.position:
            path = []
            while current_node:
                path.append(current_node.position)
                current_node = current_node.parent
            return path[::-1]

        closed_set.add(current_node.position)

        for direction in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
            neighbor_pos = (current_node.position[0] + direction[0],
                            current_node.position[1] + direction[1])

            if (0 <= neighbor_pos[0] < grid_map.shape[0] and
                0 <= neighbor_pos[1] < grid_map.shape[1] and
                grid_map[neighbor_pos] == 0 and  
                neighbor_pos not in closed_set):

                neighbor_node = Node(neighbor_pos, current_node)
                neighbor_node.g = current_node.g + 1
                neighbor_node.h = heuristic(neighbor_pos, end_node.position)
                neighbor_node.f = neighbor_node.g + neighbor_node.h

                if any(open_node.position == neighbor_node.position and open_node.f <= neighbor_node.f for open_node in open_list):
                    continue

                heapq.heappush(open_list, neighbor_node)

    return None  

def map_callback(msg):
    """订阅 `map_server` 话题，获取地图数据"""
    global map_data, map_width, map_height, resolution, origin_x, origin_y

    map_width = msg.info.width
    map_height = msg.info.height
    resolution = msg.info.resolution  # 地图分辨率（米/像素）
    origin_x = msg.info.origin.position.x  # 地图原点X坐标
    origin_y = msg.info.origin.position.y  # 地图原点Y坐标

    map_array = np.array(msg.data).reshape((map_height, map_width))
    map_data = np.where(map_array > 50, 1, 0)  # 转换为二值地图（1=障碍物，0=可通行）

    rospy.loginfo("地图已加载: {}x{}, 分辨率: {:.3f}, 原点: ({:.3f}, {:.3f})".format(
        map_width, map_height, resolution, origin_x, origin_y))

def world_to_map(x, y):
    """将世界坐标 (m) 转换为栅格坐标 (pixel)"""
    map_x = int((x - origin_x) / resolution)
    map_y = int((y - origin_y) / resolution)

    # 限制索引不超出地图范围
    map_x = np.clip(map_x, 0, map_width - 1)
    map_y = np.clip(map_y, 0, map_height - 1)

    return map_y, map_x  # 注意: ROS map 采用 (y, x) 索引顺序

def map_to_world(map_x, map_y):
    """将栅格坐标 (pixel) 转换回世界坐标 (m)"""
    world_x = map_x * resolution + origin_x
    world_y = map_y * resolution + origin_y
    return world_x, world_y

def clicked_point_callback(msg):
    """接收 Rviz 选取的目标点"""
    global start, end, map_data

    # 获取 Rviz 选取的世界坐标
    world_x = msg.pose.position.x
    world_y = msg.pose.position.y

    # 转换成栅格坐标
    grid_y, grid_x = world_to_map(world_x, world_y)

    if start is None:
        start = (grid_y, grid_x)
        rospy.loginfo("起点设定 (栅格坐标): {}".format(start))
    else:
        end = (grid_y, grid_x)
        rospy.loginfo("终点设定 (栅格坐标): {}".format(end))

        # 运行 A* 路径规划
        path = astar(map_data, start, end)

        if path:
            publish_path(path)
        else:
            rospy.logwarn("未找到可行路径！")
        start, end = None, None  # 重置起点终点

def publish_path(path):
    """发布 A* 规划的路径"""
    global path_pub

    ros_path = Path()
    ros_path.header.frame_id = "map"

    for point in path:
        pose = PoseStamped()
        pose.header.frame_id = "map"

        # 反向转换：栅格坐标 -> 世界坐标
        world_x, world_y = map_to_world(point[1], point[0])
        pose.pose.position.x = world_x
        pose.pose.position.y = world_y

        ros_path.poses.append(pose)

    path_pub.publish(ros_path)
    rospy.loginfo("A* 路径已发布")

if __name__ == '__main__':
    rospy.init_node('astar_planner')

    map_data = None
    start, end = None, None
    map_width, map_height = 0, 0
    resolution, origin_x, origin_y = 0, 0, 0

    rospy.Subscriber("/map", OccupancyGrid, map_callback)
    rospy.Subscriber("/move_base_simple/goal", PoseStamped, clicked_point_callback)
    path_pub = rospy.Publisher("/astar_path", Path, queue_size=10)

    rospy.loginfo("A* 路径规划节点启动")
    rospy.spin()
